package cn.seecoder.ai.algorithm.ml.clustering

import cn.seecoder.ai.algorithm.ml.BaseClustering
import cn.seecoder.ai.enums.TrainParamEnum
import cn.seecoder.ai.model.bo.TrainParamsBO
import cn.seecoder.ai.utils.TrainParamReader.readParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.clustering.KMeans

import java.util

// K均值聚类算法-已完成
class KMeansClustering extends BaseClustering{

  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {

    val kmeans = new KMeans()
      .setFeaturesCol("indexedFeatures")
      .setK(trainParams.k)
      .setSeed(trainParams.seed)
      .setMaxIter(trainParams.maxIter)

    kmeans
  }

}
